Python: preserve model emission order in AG-UI MESSAGES_SNAPSHOT - #7239
Merged
Evan Mattson (moonbox3) merged 7 commits intoJul 30, 2026
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
Preserves the model’s emission order when constructing AG-UI MESSAGES_SNAPSHOT in the Python ag-ui package, so lead-in text, tool calls/results, and reasoning appear in the same sequence the model streamed them (fixing #7223) while keeping the separate assistant text vs tool-call message split from #3619.
Changes:
- Added
FlowState.snapshot_segmentsand streaming-time helpers to track ordered text/tool-calls/reasoning segments. - Updated
_build_messages_snapshotto emit messages based on recorded segments (with legacy fallback when no segments exist). - Added/updated tests to assert correct snapshot ordering across mixed text/tool/reasoning scenarios.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| python/packages/ag-ui/tests/ag_ui/test_run.py | Adds tests validating snapshot ordering, tool-only ID reuse, reasoning placement, and legacy fallback behavior. |
| python/packages/ag-ui/agent_framework_ag_ui/_run_common.py | Introduces snapshot_segments and records segment boundaries during streaming emit helpers. |
| python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py | Emits snapshot messages in recorded segment order via _append_segmented_snapshot_messages, with fallback to legacy behavior. |
Contributor
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||
- Preopened message ids (tool-only path) now open a text segment when the first text arrives, so their content can't drop out of the snapshot. - A tool result closes the current tool-call segment, so call A -> result A -> call B snapshots as two pairs in stream order. - emitted_call_ids only marks calls actually emitted, keeping stale segment ids eligible for the leftover fallback. - The leftover path carries its tool results too instead of dropping them.
Evan Mattson (moonbox3)
temporarily deployed
to
github-app-auth
July 23, 2026 01:54 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
Contributor
|
Yufeng He (@he-yufeng) please fix: |
Yufeng He (he-yufeng)
temporarily deployed
to
github-app-auth
July 23, 2026 10:57 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Contributor
Author
|
Fixed in 6e5528d. The set was inferred as set[Any | None] because dict.get() keeps None in the value type even with the filter; the walrus binding narrows it properly, so update() now type-checks. Runtime behavior is unchanged. |
Yufeng He (he-yufeng)
temporarily deployed
to
github-app-auth
July 29, 2026 09:29 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Evan Mattson (moonbox3)
temporarily deployed
to
github-app-auth
July 30, 2026 10:11 — with
GitHub Actions
Inactive
GitHub Actions
Inactive
Evan Mattson (moonbox3)
approved these changes
Jul 30, 2026
Evan Mattson (moonbox3)
enabled auto-merge
July 30, 2026 10:18
Eduard van Valkenburg (eavanvalkenburg)
approved these changes
Jul 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #7223.
Problem
_build_messages_snapshotalways emitted a turn's assistant text message after its tool-call message and tool results, regardless of the order the model produced them. Lead-in narration ("Let me research...") therefore rendered below the tool chips it introduces, and post-tool summaries only looked right by accident. The rc8reasoning_messagesaccumulator had the same fixed-position issue.Approach
Track segment boundaries as the run streams.
FlowStategains asnapshot_segmentslist that the emit helpers append to: a text segment per opened text message, a tool-call segment per run of consecutive calls (shared by regular, MCP, andconfirm_changesentries), and a reasoning segment per reasoning block._build_messages_snapshotnow emits those segments in recorded order, with tool results grouped after the tool-call message they answer.The separate text/tool-call message split from #3619 is unchanged, and a tool-only turn keeps reusing the streamed message id. Flows that never went through the emit helpers fall back to the previous fixed layout, so nothing depending on it breaks.
Verification
pytest packages/ag-ui/tests— 829 passed. (test_endpoint.py/test_handoff_replay.pyfail collection on missing optionalagent_framework_orchestrations/fastapideps; identical on the base commit.)test_snapshot_reasoning_orderingencoded the fixed-position layout this issue reports; updated it to the stream-faithful order.